package ai.langchain;

import com.fasterxml.jackson.databind.JsonNode;
import com.fasterxml.jackson.databind.ObjectMapper;
import java.io.*;
import java.net.HttpURLConnection;
import java.net.URL;
import java.nio.charset.StandardCharsets;
import java.util.Arrays;
import java.util.List;

/**
 * 能跑起来，但返回的结果比给出的长多了
 */
public class OllamaRAGExample {

    private static final String OLLAMA_API_URL = "http://localhost:11434/api/generate";
    private static final String MODEL_NAME = "deepseek-r1:1.5b";
    private static final int MAX_TOKENS = 512;
    private static final ObjectMapper mapper = new ObjectMapper();

    public static void main(String[] args) {
        String question = "量子计算在密码学中有哪些影响？";
        
        try {
            // 1. 知识检索（示例使用静态数据）
            List<String> contexts = retrieveKnowledge(question);
            
            // 2. 构建增强提示
            String prompt = buildRAGPrompt(question, contexts);
            
            // 3. 调用 Ollama 模型
            String answer = callOllama(prompt);
            
            System.out.println("问题: " + question);
            System.out.println("回答: " + answer);
        } catch (Exception e) {
            e.printStackTrace();
        }
    }

    // 知识检索（模拟实现）
    private static List<String> retrieveKnowledge(String query) {
        return Arrays.asList(
            "量子计算机可能破解 RSA 加密算法，",
            "因为 Shor 算法可以在多项式时间内分解大质数，",
            "这促使了抗量子密码学的发展。"
        );
    }

    // 构建 RAG 提示模板
    private static String buildRAGPrompt(String question, List<String> contexts) {
        StringBuilder sb = new StringBuilder();
        sb.append("基于以下上下文回答：\n");
        contexts.forEach(ctx -> sb.append("- ").append(ctx).append("\n"));
        sb.append("\n问题：").append(question);
        return sb.toString();
    }

    // 调用 Ollama API
    private static String callOllama(String prompt) throws IOException {
        // 构建 Ollama 专用请求体
        String requestJson = String.format(
            "{ \"prompt\": \"%s\", \"model\": \"%s\", \"max_tokens\": %d, \"stream\": false }",
            escapeJson(prompt), MODEL_NAME, MAX_TOKENS
        );

        HttpURLConnection conn = (HttpURLConnection) new URL(OLLAMA_API_URL).openConnection();
        conn.setRequestMethod("POST");
        conn.setRequestProperty("Content-Type", "application/json");
        conn.setDoOutput(true);

        try (OutputStream os = conn.getOutputStream()) {
            os.write(requestJson.getBytes(StandardCharsets.UTF_8));
        }

        if (conn.getResponseCode() == 200) {
            try (InputStream is = conn.getInputStream()) {
                JsonNode response = mapper.readTree(is);
                return response.path("response").asText();
            }
        } else {
            throw new IOException("Ollama API 请求失败，状态码: " + conn.getResponseCode());
        }
    }

    // JSON 特殊字符转义
    private static String escapeJson(String str) {
        return str.replace("\\", "\\\\")
                .replace("\"", "\\\"")
                .replace("\n", "\\n");
    }
}